#include "FaceRecognition_DS_Final/alignment.h"

#include <math.h>
#include <stdio.h>
#include <stdlib.h>

#define SIGN(a, b) ((b) >= 0.0 ? fabs(a) : -fabs(a))
#define MAX(a, b) (a > b ? a : b)
#define MIN(a, b) (a < b ? a : b)
#define SWAP(a, b) \
    do             \
    {              \
        float tmp; \
        tmp = a;   \
        a = b;     \
        b = tmp;   \
    } while (0)

static double PYTHAG(double a, double b)
{
    double at = fabs(a), bt = fabs(b), ct, result;

    if (at > bt)
    {
        ct = bt / at;
        result = at * sqrt(1.0 + ct * ct);
    }
    else if (bt > 0.0)
    {
        ct = at / bt;
        result = bt * sqrt(1.0 + ct * ct);
    }
    else
        result = 0.0;
    return (result);
}

int DSVD(float a[][2], int m, int n, float *w, float v[][2])
{
    int flag, i, its, j, jj, k, l, nm;
    double c, f, h, s, x, y, z;
    double anorm = 0.0, g = 0.0, scale = 0.0;
    double *rv1;
    double rv1_buf[8];

    if (m < n)
    {
        fprintf(stderr, "#rows must be > #cols \n");
        return (0);
    }

    // rv1 = (double *)malloc((unsigned int) n*sizeof(double));
    rv1 = rv1_buf;

    /* Householder reduction to bidiagonal form */
    for (i = 0; i < n; i++)
    {
        /* left-hand reduction */
        l = i + 1;
        rv1[i] = scale * g;
        g = s = scale = 0.0;
        if (i < m)
        {
            for (k = i; k < m; k++)
                scale += fabs((double)a[k][i]);
            if (scale)
            {
                for (k = i; k < m; k++)
                {
                    a[k][i] = (float)((double)a[k][i] / scale);
                    s += ((double)a[k][i] * (double)a[k][i]);
                }
                f = (double)a[i][i];
                g = -SIGN(sqrt(s), f);
                h = f * g - s;
                a[i][i] = (float)(f - g);
                if (i != n - 1)
                {
                    for (j = l; j < n; j++)
                    {
                        for (s = 0.0, k = i; k < m; k++)
                            s += ((double)a[k][i] * (double)a[k][j]);
                        f = s / h;
                        for (k = i; k < m; k++)
                            a[k][j] += (float)(f * (double)a[k][i]);
                    }
                }
                for (k = i; k < m; k++)
                    a[k][i] = (float)((double)a[k][i] * scale);
            }
        }
        w[i] = (float)(scale * g);

        /* right-hand reduction */
        g = s = scale = 0.0;
        if (i < m && i != n - 1)
        {
            for (k = l; k < n; k++)
                scale += fabs((double)a[i][k]);
            if (scale)
            {
                for (k = l; k < n; k++)
                {
                    a[i][k] = (float)((double)a[i][k] / scale);
                    s += ((double)a[i][k] * (double)a[i][k]);
                }
                f = (double)a[i][l];
                g = -SIGN(sqrt(s), f);
                h = f * g - s;
                a[i][l] = (float)(f - g);
                for (k = l; k < n; k++)
                    rv1[k] = (double)a[i][k] / h;
                if (i != m - 1)
                {
                    for (j = l; j < m; j++)
                    {
                        for (s = 0.0, k = l; k < n; k++)
                            s += ((double)a[j][k] * (double)a[i][k]);
                        for (k = l; k < n; k++)
                            a[j][k] += (float)(s * rv1[k]);
                    }
                }
                for (k = l; k < n; k++)
                    a[i][k] = (float)((double)a[i][k] * scale);
            }
        }
        anorm = MAX(anorm, (fabs((double)w[i]) + fabs(rv1[i])));
    }

    /* accumulate the right-hand transformation */
    for (i = n - 1; i >= 0; i--)
    {
        if (i < n - 1)
        {
            if (g)
            {
                for (j = l; j < n; j++)
                    v[j][i] = (float)(((double)a[i][j] / (double)a[i][l]) / g);
                /* double division to avoid underflow */
                for (j = l; j < n; j++)
                {
                    for (s = 0.0, k = l; k < n; k++)
                        s += ((double)a[i][k] * (double)v[k][j]);
                    for (k = l; k < n; k++)
                        v[k][j] += (float)(s * (double)v[k][i]);
                }
            }
            for (j = l; j < n; j++)
                v[i][j] = v[j][i] = 0.0;
        }
        v[i][i] = 1.0;
        g = rv1[i];
        l = i;
    }

    /* accumulate the left-hand transformation */
    for (i = n - 1; i >= 0; i--)
    {
        l = i + 1;
        g = (double)w[i];
        if (i < n - 1)
            for (j = l; j < n; j++)
                a[i][j] = 0.0;
        if (g)
        {
            g = 1.0 / g;
            if (i != n - 1)
            {
                for (j = l; j < n; j++)
                {
                    for (s = 0.0, k = l; k < m; k++)
                        s += ((double)a[k][i] * (double)a[k][j]);
                    f = (s / (double)a[i][i]) * g;
                    for (k = i; k < m; k++)
                        a[k][j] += (float)(f * (double)a[k][i]);
                }
            }
            for (j = i; j < m; j++)
                a[j][i] = (float)((double)a[j][i] * g);
        }
        else
        {
            for (j = i; j < m; j++)
                a[j][i] = 0.0;
        }
        ++a[i][i];
    }

    /* diagonalize the bidiagonal form */
    for (k = n - 1; k >= 0; k--)
    { /* loop over singular values */
        for (its = 0; its < 30; its++)
        { /* loop over allowed iterations */
            flag = 1;
            for (l = k; l >= 0; l--)
            { /* test for splitting */
                nm = l - 1;
                if (fabs(rv1[l]) + anorm == anorm)
                {
                    flag = 0;
                    break;
                }
                if (fabs((double)w[nm]) + anorm == anorm)
                    break;
            }
            if (flag)
            {
                c = 0.0;
                s = 1.0;
                for (i = l; i <= k; i++)
                {
                    f = s * rv1[i];
                    if (fabs(f) + anorm != anorm)
                    {
                        g = (double)w[i];
                        h = PYTHAG(f, g);
                        w[i] = (float)h;
                        h = 1.0 / h;
                        c = g * h;
                        s = (-f * h);
                        for (j = 0; j < m; j++)
                        {
                            y = (double)a[j][nm];
                            z = (double)a[j][i];
                            a[j][nm] = (float)(y * c + z * s);
                            a[j][i] = (float)(z * c - y * s);
                        }
                    }
                }
            }
            z = (double)w[k];
            if (l == k)
            { /* convergence */
                if (z < 0.0)
                { /* make singular value nonnegative */
                    w[k] = (float)(-z);
                    for (j = 0; j < n; j++)
                        v[j][k] = (-v[j][k]);
                }
                break;
            }
            if (its >= 30)
            {
                // free((void*) rv1);
                fprintf(stderr, "No convergence after 30,000! iterations \n");
                return (0);
            }

            /* shift from bottom 2 x 2 minor */
            x = (double)w[l];
            nm = k - 1;
            y = (double)w[nm];
            g = rv1[nm];
            h = rv1[k];
            f = ((y - z) * (y + z) + (g - h) * (g + h)) / (2.0 * h * y);
            g = PYTHAG(f, 1.0);
            f = ((x - z) * (x + z) + h * ((y / (f + SIGN(g, f))) - h)) / x;

            /* next QR transformation */
            c = s = 1.0;
            for (j = l; j <= nm; j++)
            {
                i = j + 1;
                g = rv1[i];
                y = (double)w[i];
                h = s * g;
                g = c * g;
                z = PYTHAG(f, h);
                rv1[j] = z;
                c = f / z;
                s = h / z;
                f = x * c + g * s;
                g = g * c - x * s;
                h = y * s;
                y = y * c;
                for (jj = 0; jj < n; jj++)
                {
                    x = (double)v[jj][j];
                    z = (double)v[jj][i];
                    v[jj][j] = (float)(x * c + z * s);
                    v[jj][i] = (float)(z * c - x * s);
                }
                z = PYTHAG(f, h);
                w[j] = (float)z;
                if (z)
                {
                    z = 1.0 / z;
                    c = f * z;
                    s = h * z;
                }
                f = (c * g) + (s * y);
                x = (c * y) - (s * g);
                for (jj = 0; jj < m; jj++)
                {
                    y = (double)a[jj][j];
                    z = (double)a[jj][i];
                    a[jj][j] = (float)(y * c + z * s);
                    a[jj][i] = (float)(z * c - y * s);
                }
            }
            rv1[l] = 0.0;
            rv1[k] = f;
            w[k] = (float)x;
        }
    }
    // free((void*) rv1);
    return (1);
}

static float FrobNorm(float *f, int n)
{
    float sum = 0;

    for (int i = 0; i < n; i++)
        sum += f[i] * f[i];

    return sqrt(sum);
}

static void MatrixDot(float *a, float *b, float *c, int m, int n, int k)
{
    for (int i = 0; i < m; i++)
        for (int j = 0; j < n; j++)
        {
            c[i * n + j] = 0;
            for (int l = 0; l < k; l++)
            {
                c[i * n + j] += a[i * k + l] * b[l * n + j];
            }
        }
}

static void ComputeAffineMatrix(float cov[][2], float sigma,
                                float trans_m[][2])
{

    float u[2][2];
    float w[2];
    float v[2][2];
    float det = cov[0][0] * cov[1][1] - cov[0][1] * cov[1][0];
    float s[2][2] = {{1, 0}, {0, 1}};
    float c = 1.0;
    float r0[2][2];
    float r[2][2];

    u[0][0] = cov[0][0];
    u[0][1] = cov[0][1];

    u[1][0] = cov[1][0];
    u[1][1] = cov[1][1];

    DSVD(u, 2, 2, (float *)w, v);

    SWAP(u[0][0], u[0][1]);
    SWAP(u[1][0], u[1][1]);
    SWAP(w[0], w[1]);
    SWAP(v[0][0], v[0][1]);
    SWAP(v[1][0], v[1][1]);

    if (det < 0)
    {
        if (w[0] < w[1])
            s[1][1] = -1;
        else
            s[0][0] = -1;
    }

    MatrixDot((float *)u, (float *)s, (float *)r0, 2, 2, 2);
    MatrixDot((float *)r0, (float *)v, (float *)r, 2, 2, 2);

    if (sigma)
    {
        float diag[2][2];
        float trace[2][2];

        diag[0][0] = w[0];
        diag[0][1] = 0;
        diag[1][0] = 0;
        diag[1][1] = w[1];

        MatrixDot((float *)diag, (float *)s, (float *)trace, 2, 2, 2);

        c = 1.0 / sigma * (trace[0][0] + trace[1][1]);
    }

    trans_m[0][0] = r[0][0] * c;
    trans_m[0][1] = r[0][1] * c;
    trans_m[1][0] = r[1][0] * c;
    trans_m[1][1] = r[1][1] * c;
}

static int GetProbeVec(float *landmark, int landmark_number, float *probe_vec,
                       int probe_size, int desired_size)
{
    int i;

    float mean_face_x[] = {0.224152, 0.75610125, 0.490127, 0.254149, 0.726104};
    float mean_face_y[] = {0.2119465, 0.2119465, 0.628106, 0.780233, 0.780233};
    float from_mean[2] = {0, 0};
    float to_mean[2] = {0, 0};
    float sigma_from = 0;
    float sigma_to = 0;
    float cov[2][2] = {{0, 0}, {0, 0}};
    float trans_m[2][2];
    float padding = 0.37;

    if (landmark_number != 5 || probe_size != 2)
        return -1;

    for (i = 0; i < 5; i++)
    {
        mean_face_x[i] =
            (padding + mean_face_x[i]) / (2 * padding + 1) * desired_size;
        mean_face_y[i] =
            (padding + mean_face_y[i]) / (2 * padding + 1) * desired_size;
    }

    for (i = 0; i < 5; i++)
    {
        from_mean[0] += landmark[i];
        from_mean[1] += landmark[i + 5];
        to_mean[0] += mean_face_x[i];
        to_mean[1] += mean_face_y[i];
    }

    from_mean[0] = from_mean[0] / 5;
    from_mean[1] = from_mean[1] / 5;

    to_mean[0] = to_mean[0] / 5;
    to_mean[1] = to_mean[1] / 5;

    for (i = 0; i < 5; i++)
    {
        float gap_from[2];
        float gap_to[2];
        float distance;

        gap_from[0] = landmark[i] - from_mean[0];
        gap_from[1] = landmark[i + 5] - from_mean[1];

        distance = FrobNorm(gap_from, 2);

        sigma_from += distance * distance;

        gap_to[0] = mean_face_x[i] - to_mean[0];
        gap_to[1] = mean_face_y[i] - to_mean[1];

        distance = FrobNorm(gap_to, 2);

        sigma_to += distance * distance;

        cov[0][0] += gap_to[0] * gap_from[0];
        cov[0][1] += gap_to[0] * gap_from[1];
        cov[1][0] += gap_to[1] * gap_from[0];
        cov[1][1] += gap_to[1] * gap_from[1];
    }

    sigma_from = sigma_from / 5;
    sigma_to = sigma_to / 5;

    cov[0][0] /= 5;
    cov[0][1] /= 5;
    cov[1][0] /= 5;
    cov[1][1] /= 5;

    ComputeAffineMatrix(cov, sigma_from, trans_m);

    probe_vec[0] = trans_m[0][0];
    probe_vec[1] = trans_m[1][0];

    return 0;
}

int CalScaleAndAngle(float *landmark, int landmark_number, int desired_size,
                     float *scale, float *angle)
{
    float probe_vec[2];

    if (GetProbeVec(landmark, landmark_number, probe_vec, 2, desired_size) < 0)
        return -1;

    scale[0] = FrobNorm(probe_vec, 2);
    angle[0] = 180.0 / M_PI * atan2(probe_vec[1], probe_vec[0]);

    return 0;
}

int GetAlignedFace(const cv::Mat &img, float *landmark, int landmark_number,
                   int desired_size, cv::Mat &out)
{

    float scale;
    float angle;
    float from_center[2];
    float to_center[2];

    if (CalScaleAndAngle(landmark, landmark_number, desired_size, &scale,
                         &angle) < 0)
    {
        return -1;
    }

    to_center[0] = desired_size * 0.4;
    to_center[1] = desired_size * 0.5;

    from_center[0] = (landmark[0] + landmark[1]) / 2;
    from_center[1] = (landmark[5] + landmark[6]) / 2;

    cv::Mat rot_mat = cv::getRotationMatrix2D(
        cv::Point2f(from_center[0], from_center[1]), -1 * angle, scale);

    float ex = to_center[0] - from_center[0];
    float ey = to_center[1] - from_center[1];

    rot_mat.at<double>(0, 2) += ex;
    rot_mat.at<double>(1, 2) += ey;

    cv::warpAffine(img, out, rot_mat, cv::Size(desired_size, desired_size));
    rot_mat.release();

    return 0;
}
